import os
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import torch

def selectDataset(dataset_name,dataset_dir):
    try:
        if dataset_name == 'ucm':
            train_dir = os.path.join(dataset_dir, 'train')
            test_dir = os.path.join(dataset_dir, 'val')

            transform_train = transforms.Compose([
                transforms.Resize((256, 256)),
                transforms.RandomCrop((224, 224)),  # 224
                transforms.RandomHorizontalFlip(),
                transforms.RandomVerticalFlip(),
                #             transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2),
                transforms.ToTensor(),
                transforms.Normalize((0.485, 0.456, 0.406),
                                     (0.229, 0.224, 0.225)),
            ])
            transform_test = transforms.Compose([
                transforms.Resize((256, 256)),  # 256
                transforms.CenterCrop((224, 224)),  # 224
                transforms.ToTensor(),
                transforms.Normalize((0.485, 0.456, 0.406),
                                     (0.229, 0.224, 0.225)),
            ])

            train_set = datasets.ImageFolder(
                train_dir,
                transform_train)
            train_loader = torch.utils.data.DataLoader(
                train_set,
                batch_size=64,
                shuffle=True,
                num_workers=0,
                drop_last=True)

            test_set = datasets.ImageFolder(
                test_dir,
                transform_test)
            test_loader = torch.utils.data.DataLoader(
                test_set,
                batch_size=64,
                shuffle=False,
                num_workers=0)


        elif dataset_name == 'nwpu':
            train_dir = os.path.join(dataset_dir, 'train')
            test_dir = os.path.join(dataset_dir, 'val')
            transform_train = transforms.Compose([
                transforms.Resize((256, 256)),
                transforms.RandomCrop(256, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.RandomVerticalFlip(),
                transforms.ToTensor(),
                transforms.Normalize((0.485, 0.456, 0.406),
                                     (0.229, 0.224, 0.225)),
            ])

            transform_test = transforms.Compose([
                transforms.Resize((256, 256)),  # 256
                #             transforms.CenterCrop((224,224)),  # 224
                transforms.ToTensor(),
                transforms.Normalize((0.485, 0.456, 0.406),
                                     (0.229, 0.224, 0.225)),
            ])
            train_set = datasets.ImageFolder(
                train_dir,
                transform_train)
            train_loader = torch.utils.data.DataLoader(
                train_set,
                batch_size=64,
                shuffle=True,
                num_workers=0,
                drop_last=True)
            test_set = datasets.ImageFolder(
                test_dir,
                transform_test)
            test_loader = torch.utils.data.DataLoader(
                test_set,
                batch_size=64,
                shuffle=False,
                num_workers=0)

        elif dataset_name == 'optimal':
            train_dir = os.path.join(dataset_dir, 'train')
            test_dir = os.path.join(dataset_dir, 'val')
            transform_train = transforms.Compose([
                transforms.Resize((256, 256)),
                transforms.RandomCrop(256, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.RandomVerticalFlip(),
                transforms.ToTensor(),
                transforms.Normalize((0.485, 0.456, 0.406),
                                     (0.229, 0.224, 0.225)),
            ])

            transform_test = transforms.Compose([
                transforms.Resize((256, 256)),  # 256
                #             transforms.CenterCrop((224,224)),  # 224
                transforms.ToTensor(),
                transforms.Normalize((0.485, 0.456, 0.406),
                                     (0.229, 0.224, 0.225)),
            ])
            train_set = datasets.ImageFolder(
                train_dir,
                transform_train)
            train_loader = torch.utils.data.DataLoader(
                train_set,
                batch_size=32,
                shuffle=True,
                num_workers=0,
                drop_last=False)
            test_set = datasets.ImageFolder(
                test_dir,
                transform_test)
            test_loader = torch.utils.data.DataLoader(
                test_set,
                batch_size=64,
                shuffle=False,
                num_workers=0)

        elif dataset_name == 'aid':
            train_dir = os.path.join(dataset_dir, 'train')
            test_dir = os.path.join(dataset_dir, 'val')
            transform_train = transforms.Compose([
                transforms.RandomCrop(600, padding=10),
                transforms.RandomHorizontalFlip(),
                transforms.RandomVerticalFlip(),
                transforms.ToTensor(),
                transforms.Normalize((0.485, 0.456, 0.406),
                                     (0.229, 0.224, 0.225)),
            ])

            transform_test = transforms.Compose([
                transforms.Resize(600),  # 256
                #             transforms.CenterCrop(560),  # 224
                transforms.ToTensor(),
                transforms.Normalize((0.485, 0.456, 0.406),
                                     (0.229, 0.224, 0.225)),
            ])
            train_set = datasets.ImageFolder(
                train_dir,
                transform_train)
            train_loader = torch.utils.data.DataLoader(
                train_set,
                batch_size=8,
                shuffle=True,
                num_workers=0,
                drop_last=True)
            test_set = datasets.ImageFolder(
                test_dir,
                transform_test)
            test_loader = torch.utils.data.DataLoader(
                test_set,
                batch_size=8,
                shuffle=False,
                num_workers=0, drop_last=True)
        return train_loader, test_loader
    except:
        print("error in {} or {}".format(dataset_name,dataset_dir))

